﻿using System;
using System.Collections.Generic;
using System.IO;
using System.IO.MemoryMappedFiles;
using System.Linq;
using LibOrbisPkg.Util;

namespace LibOrbisPkg.PFS
{
  /// <summary>
  /// Class allowing parallel readonly access to a PFS archive
  /// </summary>
  public class PfsReader
  {
    /// <summary>
    /// Represents a file or directory in a PFS image.
    /// </summary>
    public abstract class Node
    {
      public Dir parent;
      public string name;
      public long offset;
      public long size;
      public long compressed_size;
      public uint ino;
      public string FullName => parent != null ? parent.FullName + "/" + name : name;
    }
    /// <summary>
    /// Represents a directory in a PFS image.
    /// </summary>
    public class Dir : Node
    {
      public List<Node> children = new List<Node>();
      public Node Get(string name)
        => children.Where(x => x.name == name).FirstOrDefault();
      public Node GetPath(string name)
      {
        var breadcrumbs = name.Split('/');
        Node n = this;
        var bc = 0;
        while(n != null && bc < breadcrumbs.Length)
        {
          n = (n as Dir)?.Get(breadcrumbs[bc]);
          bc++;
        }
        if(bc < breadcrumbs.Length)
        {
          return null;
        }
        return n;
      }
      public IEnumerable<File> GetAllFiles()
      {
        foreach(var n in children)
        {
          if (n is File f) yield return f;
          if (n is Dir d)
            foreach (var x in d.GetAllFiles())
              yield return x;
        }
      }
    }
    /// <summary>
    /// Represents a file in a PFS image.
    /// </summary>
    public class File : Node
    {
      public int[] blocks;
      private IMemoryReader reader;
      public File(IMemoryReader r) { reader = r; }
      public IMemoryReader GetView()
      {
        if (blocks != null)
          return new ChunkedMemoryReader(reader, 0x10000, blocks);
        return new MemoryAccessor(reader, offset);
      }
      public void Save(string path)
      {
        var buf = new byte[0x10000];
        using (var file = System.IO.File.OpenWrite(path))
        {
          var sz = size;
          file.SetLength(sz);
          long pos = 0;
          var reader = GetView();
          while (sz > 0)
          {
            var toRead = (int)Math.Min(sz, buf.Length);
            reader.Read(pos, buf, 0, toRead);
            file.Write(buf, 0, toRead);
            pos += toRead;
            sz -= toRead;
          }
        }
      }
    }

    // Private state for the PfsReader class
    private IMemoryReader reader;
    private PfsHeader hdr;
    private inode[] dinodes;
    private Dir root;
    private Dir uroot;
    private byte[] sectorBuf;
    private Stream sectorStream;

    public PfsReader(MemoryMappedViewAccessor r, ulong pfs_flags = 0, byte[] ekpfs = null)
    : this(new MemoryMappedViewAccessor_(r), pfs_flags, ekpfs)
    { }
    public PfsReader(IMemoryReader r, ulong pfs_flags = 0, byte[] ekpfs = null)
    {
      reader = r;
      var buf = new byte[0x400];
      reader.Read(0, buf, 0, 0x400);

      using (var ms = new MemoryStream(buf))
      {
        hdr = PfsHeader.ReadFromStream(ms);
      }
      int dinodeSize;
      Func<Stream, inode> dinodeReader;
      if (hdr.Mode.HasFlag(PfsMode.Signed))
      {
        dinodes = new DinodeS32[hdr.DinodeCount];
        dinodeReader = DinodeS32.ReadFromStream;
        dinodeSize = 0x2C8;
      }
      else
      {
        dinodes = new DinodeD32[hdr.DinodeCount];
        dinodeReader = DinodeD32.ReadFromStream;
        dinodeSize = 0xA8;
      }
      if (hdr.Mode.HasFlag(PfsMode.Encrypted))
      {
        /*if (ekpfs == null)
          throw new ArgumentException("PFS image is encrypted but no EKPFS was provided");
        var (tweakKey, dataKey) = Crypto.PfsGenEncKey(ekpfs, hdr.Seed, (pfs_flags & 0x2000000000000000UL) != 0);
		*/
		var tweakKey = new byte[] {0x7C,0xB7,0xC4,0xAC,0x8A,0x98,0x85,0x56,0x86,0xCC,0x3F,0x79,0xF7,0xBA,0x7D,0x97};
		var dataKey = new byte[] {0x04, 0xD4, 0xDE ,0xB8,0x01,0xB8,0xFF,0x7D,0x75,0x5E,0xC5,0xD9,0xDF,0x82,0xD4,0x4E};
        reader = new XtsDecryptReader(reader, dataKey, tweakKey, 16, 0x1000);
      }
      var total = 0;

      var maxPerSector = hdr.BlockSize / dinodeSize;
      sectorBuf = new byte[hdr.BlockSize];
      sectorStream = new MemoryStream(sectorBuf);
      for (var i = 0; i < hdr.DinodeBlockCount; i++)
      {
        var position = hdr.BlockSize + hdr.BlockSize * i;
        reader.Read(position, sectorBuf, 0, sectorBuf.Length);
        sectorStream.Position = 0;
        for (var j = 0; j < maxPerSector && total < hdr.DinodeCount; j++)
          dinodes[total++] = dinodeReader(sectorStream);
      }
      root = LoadDir(0, null, "");
      uroot = root.Get("uroot") as Dir;
      if (uroot == null)
        throw new Exception("Invalid PFS image (no uroot)");
      uroot.name = "uroot";
    }

    public PfsHeader Header => hdr;

    public File GetFile(string fullPath)
    {
      return uroot.GetPath(fullPath) as File;
    }

    public IEnumerable<File> GetAllFiles()
    {
      return uroot.GetAllFiles();
    }

    public Dir GetURoot()
    {
      return uroot;
    }

    public Dir GetSuperRoot()
    {
      return root;
    }

    private Dir LoadDir(uint dinode, Dir parent, string name)
    {
      var ret = new Dir() { name = name, parent = parent };
      var ino = dinodes[dinode];
      var postLoad = new List<Func<Dir>>();
      foreach (var x in Enumerable.Range(ino.StartBlock, (int)ino.Blocks))
      {
        var position = hdr.BlockSize * x;
        reader.Read(position, sectorBuf, 0, sectorBuf.Length);
        sectorStream.Position = 0;
        while (position < hdr.BlockSize * (x + 1))
        {
          var dirent = PfsDirent.ReadFromStream(sectorStream);
          if (dirent.EntSize == 0) break;
          switch (dirent.Type)
          {
            case DirentType.File:
              ret.children.Add(LoadFile(dirent.InodeNumber, ret, dirent.Name));
              break;
            case DirentType.Directory:
              postLoad.Add(() => LoadDir(dirent.InodeNumber, ret, dirent.Name));
              break;
            case DirentType.Dot:
              break;
            case DirentType.DotDot:
              break;
            default:
              break;
          }
          position += dirent.EntSize;
        }
      }
      foreach (var p in postLoad)
      {
        ret.children.Add(p());
      }
      return ret;
    }

    private File LoadFile(uint dinode, Dir parent, string name)
    {
      int[] blocks = null;
      if(dinodes[dinode].Blocks > 1 && dinodes[dinode].DirectBlocks[1] != -1)
      {
        if (!hdr.Mode.HasFlag(PfsMode.Signed))
        {
          throw new Exception("Unsigned PFS images probably shouldn't have noncontiguous blocks");
        }
        blocks = new int[dinodes[dinode].Blocks];
        var remainingBlocks = (long)dinodes[dinode].Blocks;
        var sigsPerBlock = hdr.BlockSize / 36;
        for (int i = 0; i < 12 && i < remainingBlocks; i++)
        {
          blocks[i] = dinodes[dinode].DirectBlocks[i];
        }

        remainingBlocks -= 12;
        long blockIndexOffset = 12;
        for (int i = 0; i < remainingBlocks && i < sigsPerBlock; i++)
        {
          reader.Read(dinodes[dinode].IndirectBlocks[0] * hdr.BlockSize + (i * 36) + 32, out blocks[i + blockIndexOffset]);
        }
        remainingBlocks -= sigsPerBlock;
        blockIndexOffset += sigsPerBlock;
        for (int j = 0; j * sigsPerBlock < remainingBlocks; j++)
        {
          reader.Read(dinodes[dinode].IndirectBlocks[1] * hdr.BlockSize + (j * 36) + 32, out int indirectBlockOffset);
          for (int i = 0; i < sigsPerBlock && i + (j * sigsPerBlock) < remainingBlocks; i++)
          {
            reader.Read(indirectBlockOffset * hdr.BlockSize + (i * 36) + 32, out blocks[i + blockIndexOffset]);
          }
          blockIndexOffset += sigsPerBlock;
        }
        bool contiguous = true;
        int last = blocks[0] - 1;
        for(int i = 1; i < blocks.Length; i++)
        {
          if(blocks[i - 1] + 1 != blocks[i])
          {
            contiguous = false;
            break;
          }
        }
        if (contiguous)
          blocks = null;
      }
      return new File(reader)
      {
        name = name,
        parent = parent,
        offset = dinodes[dinode].StartBlock * hdr.BlockSize,
        size = dinodes[dinode].Size,
        compressed_size = dinodes[dinode].SizeCompressed,
        ino = dinode,
        blocks = blocks,
      };
    }
  }
}
